{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Saving an equilibrium to geqdsk (via eqdsk)\n",
    "\n",
    "In this example, we take a MAST-U-like equilibrium and show how to save it to geqdsk file via the [eqdsk](https://github.com/Fusion-Power-Plant-Framework/eqdsk) package (developed by the [Bluemira](https://github.com/Fusion-Power-Plant-Framework/bluemira) team). This is an alternative to the [FreeQDSK](https://github.com/freegs-plasma/FreeQDSK)-based approach shown in the previous example (03b) - `eqdsk` additionally understands COCOS (COordinate COnventions) and can also read/write JSON and IMAS."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As in the previous example, we run our standard equilibrium solve - making it lower single null by pushing the P6 current."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# build machine\n",
    "from freegsnke import build_machine\n",
    "tokamak = build_machine.tokamak(\n",
    "    active_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle\",\n",
    "    passive_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle\",\n",
    "    limiter_path=f\"../machine_configs/MAST-U/MAST-U_like_limiter.pickle\",\n",
    "    wall_path=f\"../machine_configs/MAST-U/MAST-U_like_wall.pickle\",\n",
    ")\n",
    "\n",
    "# initialise equilibrium object\n",
    "from freegsnke import equilibrium_update\n",
    "eq = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak,\n",
    "    Rmin=0.1, Rmax=2.0,   # radial range\n",
    "    Zmin=-2.2, Zmax=2.2,  # vertical range\n",
    "    nx=65,                # number of grid points in the radial direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    ny=129,               # number of grid points in the vertical direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    ")\n",
    "\n",
    "# initialise profile object\n",
    "from freegsnke.jtor_update import ConstrainPaxisIp\n",
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq,\n",
    "    paxis=8.1e3,\n",
    "    Ip=6.2e5,\n",
    "    fvac=0.5,\n",
    "    alpha_m=1.8,\n",
    "    alpha_n=1.2\n",
    ")\n",
    "\n",
    "# initialise solver\n",
    "from freegsnke import GSstaticsolver\n",
    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)\n",
    "\n",
    "# set coil currents\n",
    "import pickle\n",
    "with open('data/simple_diverted_currents_PaxisIp.pk', 'rb') as f:\n",
    "    current_values = pickle.load(f)\n",
    "\n",
    "for key in current_values.keys():\n",
    "    eq.tokamak[key].current = current_values[key]\n",
    "\n",
    "# change this to shift plasma vertically up (-) or down (+)\n",
    "eq.tokamak[\"P6\"].current += 500\n",
    "\n",
    "# carry out forward solve\n",
    "GSStaticSolver.solve(eq=eq,\n",
    "                     profiles=profiles,\n",
    "                     constrain=None,\n",
    "                     target_relative_tolerance=1e-9)\n",
    "\n",
    "# plot the resulting equilibrium\n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing a geqdsk file\n",
    "\n",
    "`eqdsk`'s core class is `EQDSKInterface`, a dataclass holding all of the data in a geqdsk file. It's written to disk via its own `.write(path, file_format=\"geqdsk\")` method.\n",
    "\n",
    "Its field names differ slightly from FreeQDSK's (e.g. `nx`/`ny` → `nx`/`nz`, `rleft` → `xgrid1`, `simagx`/`sibdry` → `psimag`/`psibdry`, `cpasma` → `cplasma`, `rbdry`/`zbdry` → `xbdry`/`zbdry`, etc. - see the [package docs](https://github.com/Fusion-Power-Plant-Framework/eqdsk) for the full mapping). `EQDSKInterface` also carries a handful of required coil-geometry fields (`xc`/`zc`/`dxc`/`dzc`/`Ic`/`ncoil`) that plain geqdsk files don't otherwise need - FreeGSNKE tracks coil currents separately on the `tokamak` object, so we simply pass these through empty (`ncoil=0`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from eqdsk import EQDSKInterface\n",
    "\n",
    "# grid quantities\n",
    "nx, nz = eq.nx, eq.ny\n",
    "R_1D, Z_1D = eq.R_1D, eq.Z_1D\n",
    "rmin, rmax = R_1D[0], R_1D[-1]\n",
    "zmin, zmax = Z_1D[0], Z_1D[-1]\n",
    "\n",
    "# vacuum field\n",
    "# bcentre = fvac / xcentre; xcentre is an arbitrary reference radius\n",
    "# convention: often taken as the machine geometric centre or R at magnetic axis\n",
    "xcentre = 0.5 * (rmin + rmax)         # assumed — check against your COCOS/convention\n",
    "bcentre = profiles.fvac() / xcentre   # fvac passed into ConstrainPaxisIp as 0.5\n",
    "\n",
    "# magnetic axis\n",
    "xmag, zmag = eq.magneticAxis()[0:2]\n",
    "\n",
    "# last closed flux surface (i.e. core plasma boundary)\n",
    "sep = eq.separatrix(ntheta=180)        # (R, Z) points around LCFS\n",
    "xbdry, zbdry = sep[:, 0], sep[:, 1]\n",
    "\n",
    "# wall contour\n",
    "xlim, zlim = eq.tokamak.wall.R, eq.tokamak.wall.Z\n",
    "\n",
    "# 1D normalised psi\n",
    "psi_n = eq.psiN_1D(N=nx)\n",
    "\n",
    "# build the EQDSKInterface (see eqdsk docs for full field descriptions)\n",
    "eqd = EQDSKInterface(\n",
    "    name=\"MASTU_LSN\",\n",
    "    nx=nx,\n",
    "    nz=nz,\n",
    "    xdim=rmax - rmin,\n",
    "    zdim=zmax - zmin,\n",
    "    xcentre=xcentre,\n",
    "    xgrid1=rmin,\n",
    "    zmid=0.5 * (zmin + zmax),\n",
    "    xmag=xmag,\n",
    "    zmag=zmag,\n",
    "    psimag=eq.psi_axis,\n",
    "    psibdry=eq.psi_bndry,\n",
    "    bcentre=bcentre,\n",
    "    cplasma=eq.plasmaCurrent(),\n",
    "    fpol=profiles.fpol(psi_n),\n",
    "    pressure=profiles.pressure(psi_n),\n",
    "    ffprime=profiles.ffprime(psi_n),\n",
    "    pprime=profiles.pprime(psi_n),\n",
    "    psi=eq.psi(),\n",
    "    qpsi=eq.q(psi_n),\n",
    "    nbdry=len(xbdry),\n",
    "    xbdry=xbdry,\n",
    "    zbdry=zbdry,\n",
    "    nlim=len(xlim),\n",
    "    xlim=xlim,\n",
    "    zlim=zlim,\n",
    "    # no coil geometry to store in a plain geqdsk file\n",
    "    ncoil=0,\n",
    "    xc=np.array([]),\n",
    "    zc=np.array([]),\n",
    "    dxc=np.array([]),\n",
    "    dzc=np.array([]),\n",
    "    Ic=np.array([]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the file!\n",
    "eqd.write(\"MASTU_LSN.geqdsk\", file_format=\"geqdsk\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading a geqdsk file\n",
    "\n",
    "`eqdsk` reads a geqdsk file back in via `EQDSKInterface.from_file(path)`, returning an `EQDSKInterface` instance whose fields are accessed as attributes rather than dictionary keys.\n",
    "\n",
    "By default, `from_file` also tries to *identify* the file's COCOS and convert it to a target convention (COCOS 11) - useful when combining eqdsks written by different codes, but unnecessary here since we know exactly where this file came from. We pass `no_cocos=True` to skip that step and just get the raw values straight back.\n",
    "\n",
    "A geqdsk file doesn't store individual coil currents, so - just as in the FreeQDSK example - we use the inverse solver to re-build the equilibrium, using the stored 2D psi map (and other data) as constraints.\n",
    "\n",
    "We build the constraints entirely from what the geqdsk file gives us:\n",
    "\n",
    "- an **isoflux constraint** from the last closed flux surface (`xbdry`/`zbdry`) - all these points should sit on the same flux contour, without us needing to know its absolute value;\n",
    "- **null-point constraints** at the magnetic axis (`xmag`/`zmag`, also from the file);\n",
    "- a **`psi_vals` constraint built from the *entire* saved 2D `psi` map** - unlike the constraints above, this one is given the actual flux *values*, not just which points should share a value, so it's what lets the optimiser pin down the absolute flux level (and, with it, the individual coil currents) rather than just the plasma shape.\n",
    "\n",
    "To check that the file faithfully captures the equilibrium, we'll reconstruct it purely from what's in the geqdsk file: the grid extent, the plasma current (`cplasma`), and the `p'(ψ_n)`/`FF'(ψ_n)` profile data (`pprime`/`ffprime`).\n",
    "\n",
    "FreeGSNKE's `GeneralPprimeFFprime` profile class accepts exactly this kind of tabulated profile data directly, rather than a parametric form like `ConstrainPaxisIp` above.\n",
    "\n",
    "We reset the coil currents on a **fresh, independent** tokamak to zero first, to make sure we aren't just reusing the answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# open the file!\n",
    "eqd_read = EQDSKInterface.from_file(\"MASTU_LSN.geqdsk\", no_cocos=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke.jtor_update import GeneralPprimeFFprime\n",
    "\n",
    "# grid data\n",
    "Rmin = eqd_read.xgrid1\n",
    "Rmax = Rmin + eqd_read.xdim\n",
    "Zmin = eqd_read.zmid - 0.5 * eqd_read.zdim\n",
    "Zmax = eqd_read.zmid + 0.5 * eqd_read.zdim\n",
    "nx_read, ny_read = eqd_read.nx, eqd_read.nz\n",
    "\n",
    "# normalised psi\n",
    "psi_n_read = np.linspace(0, 1, nx_read)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's set up and run the inverse solver, using those quantities plus the constraints described above to reconstruct the coil currents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke.inverse import Inverse_optimizer\n",
    "\n",
    "# coil names\n",
    "active_coil_names = tokamak.coils_list[: tokamak.n_active_coils]\n",
    "\n",
    "# fresh tokamak so we can set coil currents to zero and compare to old eq object\n",
    "tokamak_inv = build_machine.tokamak(\n",
    "    active_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle\",\n",
    "    passive_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle\",\n",
    "    limiter_path=f\"../machine_configs/MAST-U/MAST-U_like_limiter.pickle\",\n",
    "    wall_path=f\"../machine_configs/MAST-U/MAST-U_like_wall.pickle\",\n",
    ")\n",
    "\n",
    "# fresh equilibrium object\n",
    "eq_inv = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak_inv,\n",
    "    Rmin=Rmin, Rmax=Rmax,\n",
    "    Zmin=Zmin, Zmax=Zmax,\n",
    "    nx=nx_read,\n",
    "    ny=ny_read,\n",
    ")\n",
    "\n",
    "# fresh profiles object\n",
    "profiles_inv = GeneralPprimeFFprime(\n",
    "    eq=eq_inv,\n",
    "    Ip=eqd_read.cplasma,\n",
    "    fvac=eqd_read.bcentre * eqd_read.xcentre,\n",
    "    psi_n=psi_n_read,\n",
    "    pprime_data=eqd_read.pprime,\n",
    "    ffprime_data=eqd_read.ffprime,\n",
    ")\n",
    "\n",
    "# isoflux constraint: subsample the LCFS boundary stored in the geqdsk file\n",
    "isoflux_set = [[eqd_read.xbdry[::5], eqd_read.zbdry[::5]]]\n",
    "\n",
    "# null-point constraints: magnetic axis (from the file)\n",
    "null_points = [\n",
    "    [eqd_read.xmag],\n",
    "    [eqd_read.zmag],\n",
    "]\n",
    "\n",
    "# psi_vals constraint built from the *entire* saved 2D psi map. Passing the grid's own\n",
    "# (R, Z) meshgrids (rather than a sparse list of points) triggers Inverse_optimizer's\n",
    "# \"full grid\" fast path, which reuses the equilibrium's own cached Greens functions\n",
    "# instead of recomputing them pointwise.\n",
    "psi_vals = [eq_inv.R, eq_inv.Z, eqd_read.psi]\n",
    "\n",
    "# build constraints\n",
    "constrain = Inverse_optimizer(\n",
    "    isoflux_set=isoflux_set,\n",
    "    null_points=null_points,\n",
    "    psi_vals=psi_vals,\n",
    ")\n",
    "\n",
    "# initialise solver object\n",
    "GSStaticSolver_inv = GSstaticsolver.NKGSsolver(eq_inv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set coils to zero current\n",
    "for label in active_coil_names:\n",
    "    eq_inv.tokamak.set_coil_current(label, 0.0)\n",
    "\n",
    "# default plasma flux guess\n",
    "eq_inv.plasma_psi = eq_inv.create_psi_plasma_default()\n",
    "\n",
    "# solve\n",
    "GSStaticSolver_inv.solve(\n",
    "    eq=eq_inv,\n",
    "    profiles=profiles_inv,\n",
    "    constrain=constrain,\n",
    "    target_relative_tolerance=1e-4,\n",
    "    target_relative_psit_update=1e-3,\n",
    "    max_solving_iterations=50,\n",
    "    l2_reg=1e-9,\n",
    "    verbose=True,\n",
    "    Picard_handover=1e-5,\n",
    "    full_jacobian_handover=[2e-2,1e-2],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This converges properly and gives the correct total flux, coil currents close to their true values, and `psi_axis`/`psi_bndry` matching to several decimal places."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print coil currents\n",
    "reconstructed_currents = {label: eq_inv.tokamak[label].current for label in active_coil_names}\n",
    "original_active_currents = {label: eq.tokamak[label].current for label in active_coil_names}\n",
    "\n",
    "print(\"Coil currents [A]:\")\n",
    "print(f\"{'coil':>10s}  {'original':>12s}  {'reconstructed':>14s}\")\n",
    "for label in active_coil_names:\n",
    "    print(f\"{label:>10s}  {original_active_currents[label]:12.1f}  {reconstructed_currents[label]:14.1f}\")\n",
    "\n",
    "# some small differences\n",
    "psi_diff_inv = np.abs(eq.psi() - eq_inv.psi())\n",
    "print(f\"\\nMagnetic axis:   original = {eq.magneticAxis()[0:2]}, reconstructed = {eq_inv.magneticAxis()[0:2]}\")\n",
    "print(f\"psi_axis:        original = {eq.psi_axis:.6g}, reconstructed = {eq_inv.psi_axis:.6g}\")\n",
    "print(f\"psi_boundary:    original = {eq.psi_bndry:.6g}, reconstructed = {eq_inv.psi_bndry:.6g}\")\n",
    "print(f\"Max |Δpsi| relative to max|psi|: {np.max(psi_diff_inv) / np.max(np.abs(eq.psi())):.3e}\")\n",
    "\n",
    "# plot the two equilibria side by side\n",
    "fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8), dpi=80)\n",
    "for ax, e, title in [(ax1, eq, \"Original\"), (ax2, eq_inv, \"Inverse reconstruction\")]:\n",
    "    ax.grid(True, which=\"both\")\n",
    "    e.plot(axis=ax, show=False)\n",
    "    e.tokamak.plot(axis=ax, show=False)\n",
    "    ax.set_xlim(0.1, 2.15)\n",
    "    ax.set_ylim(-2.25, 2.25)\n",
    "    ax.set_title(title)\n",
    "    if ax == ax2:\n",
    "        constrain.plot(axis=ax2)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "freegsnke",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}